import os
import numpy as np
import json


def data_preprocess(datapath):
    trains = []
    fin = open(os.path.join(datapath, "train.txt"), "r", encoding="utf-8")
    for line in fin:
        elems = line.strip().split(",")
        try:
            assert len(elems) == 3
        except:
            print("error")
        trains.append(elems)

    entities = list(set([triple[0] for triple in trains]) | set([triple[2] for triple in trains]))
    relations = list(set([triple[1] for triple in trains]))

    entity_n, relation_n = len(entities), len(relations)
    info = {"n_entity": entity_n, "n_relation": relation_n}
    with open(os.path.join(datapath, "info.json"), "w", encoding="utf-8") as fout:
        json.dump(info, fout)

    ent_text2id = {entity: idx for idx, entity in enumerate(entities)}
    rel_text2id = {relation: idx for idx, relation in enumerate(relations)}

    ent_id2text = {idx: entity for idx, entity in enumerate(entities)}
    rel_id2text = {idx: relation for idx, relation in enumerate(relations)}

    trains_id = [[ent_text2id[triple[0]], rel_text2id[triple[1]], ent_text2id[triple[2]]] for triple in trains]

    with open(os.path.join(datapath, "entities.txt"), "w", encoding="utf-8") as fout:
        for entity in entities:
            fout.write(entity + "\n")

    with open(os.path.join(datapath, "relations.txt"), "w", encoding="utf-8") as fout:
        for relation in relations:
            fout.write(relation + "\n")

    with open(os.path.join(datapath, "entity2id.txt"), "w", encoding="utf-8") as fout:
        fout.write(str(len(entities)) + '\n')
        for idx, entity in enumerate(entities):
            fout.write('\t'.join([entity, str(idx)]) + '\n')

    with open(os.path.join(datapath, "relation2id.txt"), "w", encoding="utf-8") as fout:
        fout.write(str(len(relations)) + '\n')
        for idx, relation in enumerate(relations):
            fout.write('\t'.join([relation, str(idx)]) + '\n')

    with open(os.path.join(datapath, "ent_text2id.dict"), "w", encoding="utf-8") as fout:
        json.dump(ent_text2id, fout)

    with open(os.path.join(datapath, "rel_text2id.dict"), "w", encoding="utf-8") as fout:
        json.dump(rel_text2id, fout)

    with open(os.path.join(datapath, "ent_id2text.dict"), "w", encoding="utf-8") as fout:
        json.dump(ent_id2text, fout)

    with open(os.path.join(datapath, "rel_id2text.dict"), "w", encoding="utf-8") as fout:
        json.dump(rel_id2text, fout)

    np.savetxt(os.path.join(datapath, "train_id.txt"), trains_id, delimiter="\t", fmt="%d")


if __name__ == "__main__":
    data_path = os.path.join('../data/PowerGrid')
    data_preprocess(data_path)
